package main

import (
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pion/webrtc/v4"
	"github.com/pion/webrtc/v4/pkg/media"
	"github.com/pion/webrtc/v4/pkg/media/ivfwriter"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

type Connection struct {
	*websocket.Conn
	sync.Mutex
}

func main() {
	// 启动HTTP服务器
	http.HandleFunc("/", serveHome)
	http.HandleFunc("/ws", serveWS)
	go func() {
		if err := http.ListenAndServe(":8080", nil); err != nil {
			log.Fatal(err)
		}
	}()

	// 准备WebRTC配置
	config := webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{
				URLs: []string{"stun:stun.l.google.com:19302"},
			},
		},
	}

	// 创建PeerConnection
	peerConnection, err := webrtc.NewPeerConnection(config)
	if err != nil {
		log.Fatal(err)
	}

	// 添加视频轨道
	videoTrack, err := webrtc.NewTrackLocalStaticSample(
		webrtc.RTPCodecCapability{
			MimeType:    "video/vp8",
			ClockRate:   90000,
			Channels:    0,
		},
		"video",
		"pion",
	)
	if err != nil {
		log.Fatal(err)
	}
	_, err = peerConnection.AddTrack(videoTrack)
	if err != nil {
		log.Fatal(err)
	}

	// 处理ICE候选
	peerConnection.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c == nil {
			return
		}
		fmt.Printf("ICE Candidate: %s\n", c.ToJSON().Candidate)
	})

	// 启动FFmpeg命令读取视频
	cmd := exec.Command("ffmpeg", 
		"-i", "input.mp4", 
		"-an",                 // 禁用音频
		"-c:v", "libvpx",      // VP8编码
		"-f", "ivf",           // IVF容器格式
		"-t", "60",            // 限制为60秒
		"pipe:1",
	)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}

	if err = cmd.Start(); err != nil {
		log.Fatal(err)
	}

	// 创建IVF解析器
	ivfReader, header, err := ivfreader.NewWith(stdout)  // 改为ivfreader
	if err != nil {
		log.Fatal(err)
	}

	// 设置视频轨道参数
	videoTrack.UpdateCodec(webrtc.RTPCodecCapability{
		MimeType:     "video/vp8",
		ClockRate:    90000,
		Channels:     0,
		SDPFmtpLine:  "",
		RTCPFeedback: nil,
	})

	// 读取视频帧并发送
	ticker := time.NewTicker(time.Millisecond * time.Duration((1000/header.TimebaseNumerator)*header.TimebaseDenominator))
	for range ticker.C {
		frame, _, err := ivfReader.ParseNextFrame()
		if err != nil {
			log.Fatal(err)
		}

		if err = videoTrack.WriteSample(media.Sample{
			Data:    frame,
			Duration: time.Second,
		}); err != nil {
			log.Fatal(err)
		}
	}
}

func serveHome(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

func serveWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("Upgrade:", err)
		return
	}
	defer conn.Close()

	// 处理WebSocket消息
	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			log.Println("Read:", err)
			break
		}
		log.Printf("Recv: %s", message)
	}
}